# The import statement will vary depending on your LLM and vector database. This is an example for OpenAI + ChromaDB
from ast import main
from pymilvus import MilvusClient
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.milvus.milvus_vector import Milvus_VectorStore
from openai import OpenAI
from pymilvus import model
def pymilvus_embedding_function():
    sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(
        model_name='/home/tom/llms/bge-small-zh-v1.5',  # 指定模型名称
        device='cpu'  # 指定设备，例如 'cpu' 或 'cuda:0'
    )
    return sentence_transformer_ef

class MyVanna(Milvus_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):

        milvus_client = MilvusClient(uri="http://localhost:19530")
        config = {
            "embedding_function": pymilvus_embedding_function(),
            "milvus_client": milvus_client,
            "model": "deepseek-chat",
            "temperature": 0.7,
            "dialect": "SQLLite"
        }
        client = OpenAI(api_key="sk-79fa0380ce9c4c3297e51451baab0d09", base_url="https://api.deepseek.com/v1")
        Milvus_VectorStore.__init__(self, config=config)
        # config中包含模型信息 config= {"model": "deepseek-chat", "temperature": 0.7}
        OpenAI_Chat.__init__(self, client=client, config=config)
vn = MyVanna()
# vn.connect_to_sqlite("https://vanna.ai/chinook.sqlite")
vn.connect_to_sqlite("./Chinook.sqlite")
def train():
    # 训练表结构信息
    # 1、获取sqlite数据库中所有的表结构信息
    sql = "SELECT type, sql FROM sqlite_master WHERE sql is not null"
    tables = vn.run_sql(sql)
    # 2、根据表结构信息进行训练
    for ddl in tables["sql"].tolist():
        vn.train(ddl=ddl)

if __name__ == '__main__':

    # train()
    obj = vn.ask("每位顾客在各流派上花费了多少？",)
    print(obj)